import cv2
import numpy as np
# from collections import OrderedDict
from filterpy.kalman import KalmanFilter
from ultralytics import YOLO


# ----------------------------- 卡尔曼滤波器定义 -----------------------------
class KalmanFilterTracker:
    count = 0  # 用于分配唯一的ID

    def __init__(self, bbox):
        """
        初始化卡尔曼滤波器跟踪器
        :param bbox: 初始边界框 [x1, y1, x2, y2]
        """
        # 定义状态向量 [cx, cy, area, ratio, vx, vy, varea, vratio]
        # cx, cy 是中心点坐标; area是面积; ratio是宽高比; v*是对应的速度
        self.kf = KalmanFilter(dim_x=8, dim_z=4)
        self.kf.F = np.array(
            [
                [1, 0, 0, 0, 1, 0, 0, 0],  # 状态转移矩阵
                [0, 1, 0, 0, 0, 1, 0, 0],
                [0, 0, 1, 0, 0, 0, 1, 0],
                [0, 0, 0, 1, 0, 0, 0, 1],
                [0, 0, 0, 0, 1, 0, 0, 0],
                [0, 0, 0, 0, 0, 1, 0, 0],
                [0, 0, 0, 0, 0, 0, 1, 0],
                [0, 0, 0, 0, 0, 0, 0, 1]
            ],
            dtype=float)

        self.kf.H = np.array(
            [
                [1, 0, 0, 0, 0, 0, 0, 0],  # 观测矩阵
                [0, 1, 0, 0, 0, 0, 0, 0],
                [0, 0, 1, 0, 0, 0, 0, 0],
                [0, 0, 0, 1, 0, 0, 0, 0]
            ],
            dtype=float)

        self.kf.P *= 1000.  # 初始协方差矩阵
        self.kf.R *= 10  # 观测噪声
        self.kf.Q *= 0.01  # 过程噪声

        # 初始化状态
        cx = (bbox[0] + bbox[2]) / 2
        cy = (bbox[1] + bbox[3]) / 2
        w = bbox[2] - bbox[0]
        h = bbox[3] - bbox[1]
        area = w * h
        ratio = w / float(h)
        # 状态向量 [cx, cy, area, ratio, vx, vy, varea, vratio]
        self.kf.x = np.array([cx, cy, area, ratio, 0, 0, 0, 0])

        self.id = KalmanFilterTracker.count
        KalmanFilterTracker.count += 1
        self.history = []  # 存储历史轨迹点
        self.hits = 0
        self.no_losses = 0

    def update(self, bbox):
        """
        使用检测到的边界框更新卡尔曼滤波器
        :param bbox: 检测到的边界框 [x1, y1, x2, y2]
        """
        self.hits += 1
        self.no_losses = 0

        cx = (bbox[0] + bbox[2]) / 2
        cy = (bbox[1] + bbox[3]) / 2
        w = bbox[2] - bbox[0]
        h = bbox[3] - bbox[1]
        area = w * h
        ratio = w / float(h)
        measurement = np.array([cx, cy, area, ratio])
        self.kf.update(measurement)

    def predict(self):
        """
        预测下一帧的状态
        :return: 预测的边界框 [x1, y1, x2, y2]
        """
        self.kf.predict()
        pred_cx, pred_cy, pred_area, pred_ratio = self.kf.x[:4]

        # 检查是否有NaN值
        if np.any(np.isnan([pred_cx, pred_cy, pred_area, pred_ratio])):
            # 如果有NaN值，返回一个标识无效的边界框
            return [np.nan, np.nan, np.nan, np.nan]

        pred_w = np.sqrt(pred_area * pred_ratio)
        pred_h = pred_area / pred_w
        pred_bbox = [
            pred_cx - pred_w / 2, pred_cy - pred_h / 2, pred_cx + pred_w / 2,
            pred_cy + pred_h / 2
        ]

        # 更新历史轨迹
        self.history.append((int(pred_cx), int(pred_cy)))
        if len(self.history) > 50:  # 限制轨迹长度
            self.history.pop(0)

        return pred_bbox


# ----------------------------- 匈牙利算法或简单匹配 (这里使用IoU匹配) -----------------------------
def iou(bb_test, bb_gt):
    """
    计算两个边界框的IoU
    :param bb_test: 检测框 [x1, y1, x2, y2]
    :param bb_gt: 真实框或预测框 [x1, y1, x2, y2]
    :return: IoU值
    """
    xx1 = np.maximum(bb_test[0], bb_gt[0])
    yy1 = np.maximum(bb_test[1], bb_gt[1])
    xx2 = np.minimum(bb_test[2], bb_gt[2])
    yy2 = np.minimum(bb_test[3], bb_gt[3])
    w = np.maximum(0., xx2 - xx1)
    h = np.maximum(0., yy2 - yy1)
    wh = w * h
    o = wh / ((bb_test[2] - bb_test[0]) * (bb_test[3] - bb_test[1]) +
              (bb_gt[2] - bb_gt[0]) * (bb_gt[3] - bb_gt[1]) - wh)
    return o


def associate_detections_to_trackers(detections,
                                     trackers,
                                     tracker_objects,
                                     iou_threshold=0.3):
    """
    将检测结果与现有跟踪器进行匹配
    :param detections: 检测结果列表 [[x1, y1, x2, y2], ...]
    :param trackers: 跟踪器预测框列表，每个元素是边界框 [x1, y1, x2, y2]
    :param tracker_objects: KalmanFilterTracker对象列表
    :param iou_threshold: IoU匹配阈值
    :return: matched_indices, unmatched_detections, unmatched_trackers
    """
    if len(trackers) == 0:
        return np.empty(
            (0, 2), dtype=int), np.arange(len(detections)), np.empty((0, 1),
                                                                     dtype=int)

    iou_matrix = np.zeros((len(detections), len(trackers)), dtype=np.float32)
    for d, det in enumerate(detections):
        for t, trk_pred in enumerate(trackers):
            iou_matrix[d, t] = iou(det, trk_pred)  # 使用预测框计算IoU

    # 简单的匹配策略：最大IoU匹配
    matched_indices = []
    unmatched_detections = list(range(len(detections)))
    unmatched_trackers = list(range(len(trackers)))

    if min(iou_matrix.shape) > 0:
        a = (iou_matrix > iou_threshold).astype(np.int32)
        if a.sum(1).max() == 1 and a.sum(0).max() == 1:
            matched_indices = np.stack(np.where(a), axis=1)
        else:
            # 使用匈牙利算法或其他更复杂的匹配方法会更好
            # 这里简化处理
            # 修改：使用更合理的匹配逻辑
            for d_idx in range(len(detections)):
                if len(unmatched_trackers) == 0:
                    break
                # 找到当前检测与所有未匹配跟踪器的最大IoU
                max_iou = 0
                max_t_idx = -1
                for t_idx in unmatched_trackers:
                    if iou_matrix[d_idx, t_idx] > max_iou:
                        max_iou = iou_matrix[d_idx, t_idx]
                        max_t_idx = t_idx

                # 如果最大IoU超过阈值，则匹配
                if max_iou > iou_threshold:
                    matched_indices.append([d_idx, max_t_idx])
                    if d_idx in unmatched_detections:
                        unmatched_detections.remove(d_idx)
                    if max_t_idx in unmatched_trackers:
                        unmatched_trackers.remove(max_t_idx)

    matched_indices = np.array(matched_indices)
    unmatched_detections = np.array(unmatched_detections)
    unmatched_trackers = np.array(unmatched_trackers)

    return matched_indices, unmatched_detections, unmatched_trackers


# ----------------------------- 主函数 -----------------------------
def main():
    # 加载YOLOv8模型
    model = YOLO('yolov8n.pt')

    # 打开视频文件或摄像头
    # video_path = "path/to/your/video.mp4" # 替换为你的视频路径
    # cap = cv2.VideoCapture(video_path)
    cap = cv2.VideoCapture(0)  # 使用摄像头
    if not cap.isOpened():
        print("Error: Could not open video/camera.")
        return

    # 存储跟踪器
    trackers = []
    max_age = 10  # 跟踪器最大未匹配帧数，超过则删除

    while True:
        ret, frame = cap.read()
        if not ret:
            print("End of video or cannot read the frame.")
            break

        # YOLOv8检测
        results = model(frame)
        detections = []
        confidences = []
        classes = []

        # 处理检测结果 (只保留置信度高的 'person' 类别作为示例)
        for r in results:
            boxes = r.boxes
            for box in boxes:
                cls = int(box.cls[0])
                conf = float(box.conf[0])
                if cls == 0 and conf > 0.5:  # 0 是 'person' 类别的ID
                    x1, y1, x2, y2 = box.xyxy[0].cpu().numpy().astype(int)
                    detections.append([x1, y1, x2, y2])
                    confidences.append(conf)
                    classes.append(cls)

        # 转换跟踪器列表为边界框列表用于匹配
        trks = []
        to_del = []
        for t, trk in enumerate(trackers):
            pos = trk.predict()
            trks.append(pos)
            if np.any(np.isnan(pos)):
                to_del.append(t)
        # 删除无效跟踪器
        for t in reversed(to_del):
            trackers.pop(t)
        # 更新trks索引以匹配删除后的trackers
        unmatched_trks_map = [i for i in range(len(trackers))]

        # 匹配检测框和跟踪器
        matched, unmatched_dets, unmatched_trks = associate_detections_to_trackers(
            detections, trks, trackers)

        # 更新匹配上的跟踪器
        for m in matched:
            # 修复：确保访问的tracker索引在有效范围内
            if m[1] < len(trackers) and m[0] < len(detections):
                trackers[m[1]].update(detections[m[0]])

        # 为未匹配的检测创建新的跟踪器
        for i in unmatched_dets:
            # 修复：确保检测索引在有效范围内
            if i < len(detections):
                trk = KalmanFilterTracker(detections[i])
                trackers.append(trk)

        # 更新未匹配的跟踪器 (增加未匹配计数)
        for i in unmatched_trks:
            # 修复：确保索引有效
            if i < len(trackers):
                trackers[i].no_losses += 1

        # 删除长时间未匹配的跟踪器
        trackers = [t for t in trackers if t.no_losses <= max_age]

        # 绘制结果
        # 绘制检测框 (绿色)
        for det in detections:
            x1, y1, x2, y2 = map(int, det)
            cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)

        # 绘制跟踪框和轨迹 (红色)
        for trk in trackers:
            if trk.hits > 0:  # 只绘制至少被更新过一次的跟踪器
                pred_bbox = trk.predict()  # 获取预测框用于绘制
                # 检查预测框是否有效（不包含NaN）
                if not np.any(np.isnan(pred_bbox)):
                    x1, y1, x2, y2 = map(int, pred_bbox)
                    cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 0, 255), 2)
                    cv2.putText(frame, f'ID: {trk.id}', (x1, y1 - 10),
                                cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 2)
                    # 绘制轨迹
                    for i in range(1, len(trk.history)):
                        cv2.line(frame, trk.history[i - 1], trk.history[i],
                                 (255, 0, 0), 2)

        # 显示结果
        cv2.imshow('YOLOv8 + Kalman Filter Tracking', frame)

        # 按 'q' 键退出
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

    # 释放资源
    cap.release()
    cv2.destroyAllWindows()


if __name__ == "__main__":
    main()
